# @Author：cnzdy
# @Email：cn_zdy@126.com
# @Time: 2021/10/5 10:53
# @File: evaluator.py

from torchnet import meter
from tqdm import tqdm
import pandas as pd
import numpy as np
import os
import cv2
import matplotlib.pyplot as plt
from utils_analysis.metrics import plot_confusion_matrix, eval_metrics, plot_roc
import torch

from pytorch.classifier.analysis import tensor_show
from pytorch.classifier.dataset import get_loaders
from pytorch.classifier.options import Options

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


class Evaluator(object):
    """模型评估"""

    def __init__(self):
        self.model = Options.model
        self.model.to(device)
        self.model.eval()
        self.loader = Options.test_loader
        self.class_names = Options.class_names
        self.result_file = Options.test_result_file

    # 验证模型时不需要求导，即不需要梯度计算，关闭autograd，可以提高速度，节约内存。如果不关闭可能会爆显存。
    @torch.no_grad()
    def test(self):
        """在测试集上模型执行预测，将预测结果输出csv文件中"""
        results = []
        ground_truth = []
        pred_probabilities = []

        cm = meter.ConfusionMeter(2)
        loader = tqdm(self.loader)

        for ii, (inputs, labels) in enumerate(loader):
            test_input = inputs.to(device, non_blocking=True)
            test_label = labels.to(device, non_blocking=True)

            test_score = self.model(test_input)
            probabilities = torch.nn.functional.softmax(test_score, dim=1)[:, 1].detach().tolist()
            cm.add(test_score.detach().squeeze(), test_label.type(torch.LongTensor))

            ground_truth += test_label.cpu().tolist()
            pred_probabilities += probabilities
            # print(f"probabilities type: {type(probabilities)}")

        # print(f"ground_truth: {ground_truth}")
        # print(f"pred_results: {pred_probabilities}")
        self.save_results(ground_truth, pred_probabilities)

        # 计算模型的各项指标
        eval_metrics(ground_truth, pred_probabilities, cm.value(), Options.analysis_path)

        return results

    def save_results(self, ground_truth, pred_results):
        """保存预测结果"""
        results = [(label, prob) for label, prob in zip(ground_truth, pred_results)]
        # print(f"results type: {type(results)}")

        df = pd.DataFrame(data=results)
        df.to_csv(self.result_file, encoding='gbk')

    # @pysnooper.snoop()
    def predict(self, img_file):
        """预测单张图片的类别，并显示.

        Args:
            img_file:图片的绝对路径.
        """
        self.model.cpu()
        img = cv2.imread(img_file)
        # print(type(img))
        img.resize([1, 224, 224, 3])
        img = img / 255
        img = np.transpose(img, (0, 3, 1, 2))
        img = torch.from_numpy(img).float()
        output = self.model(img)
        print(output)

    @torch.no_grad()
    def viz_pred(self, num_images=6, save_path='.'):
        """选取n张图片预测，并在图片上显示预测结果
        :param num_images: 默认选取的图片为6
        """
        self.model.to(device)
        images_so_far = 0
        plt.figure()
        loader = tqdm(self.loader)

        for i, (inputs, labels) in enumerate(loader):
            inputs = inputs.to(device, non_blocking=True)
            # labels = labels.to(device)

            outputs = self.model(inputs)
            _, preds = torch.max(outputs, 1)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images // 2, 2, images_so_far)
                ax.axis('off')
                ax.set_title('predicted: {}'.format(self.class_names[preds[j]]))
                tensor_show(inputs.cpu().data[j])

                if images_so_far == num_images:
                    if save_path.strip() != '':
                        plt.savefig(os.path.join(save_path, 'predict_images.png'), dpi=200)
                    plt.show()
                    return


def eval():
    if Options.model_file.strip() != '':
        Options.model.load_state_dict(torch.load(Options.model_file))
        evaltor = Evaluator()
        evaltor.test()
        # -------------------------------------------#
        evaltor.predict(r"M:\classification\cat-dog\data\mini\test\cats\cat.200.jpg")
        evaltor.viz_pred(save_path=Options.analysis_path)


if __name__ == '__main__':
    get_loaders(True)
    eval()
